"""
    Title

    author: wxz
    date: 2021-12-11
    github: https://github.com/xinzwang
"""
from pathlib import Path

import torch
import torch.nn as nn

import sys

sys.path.append('..')

from yolov5.models.experimental import attempt_load
from yolov5.utils.torch_utils import select_device


def ModelLoadHelper(weights='yolov5/weights/yolov5s.pt', device=None, batch_size=None, dnn=True):
    device = select_device(device, batch_size)
    model = DetectBackend(weights, device, dnn)
    model.eval()
    model.warmup()
    return model, device


class DetectBackend(nn.Module):
    def __init__(self, weights='yolov5/weights/yolov5s.pt', device=None, dnn=True):
        super().__init__()
        suffix = Path(weights).suffix.lower()
        assert suffix == '.pt', 'weights must be .pt file'  # 必须是pytorch的.pt文件

        model = attempt_load(weights, map_location=device)
        stride = int(model.stride.max())
        names = model.modules().name if hasattr(model, 'module') else model.names
        self.__dict__.update(locals())  # assign all variables to self
        return

    def forward(self, im, augment=False, visualize=False):
        y = self.model(im, augment=augment, visualize=visualize)
        return y

    def warmup(self, img_size=(1, 3, 640, 640)):
        """执行一次推理    预热模型"""
        if isinstance(self.device, torch.device) and self.device.type != 'cpu':
            im = torch.zeros(*img_size).to(self.device).type(torch.float)
            self.forward(im)


if __name__ == '__main__':
    from tqdm import tqdm
    from yolo_data_loader import ImageLoaderHelper
    from yolov5.utils.torch_utils import time_sync

    path = '../../datasets/coco/my_test.txt'
    # path = '../../datasets/coco/val2017.txt'
    loader, _ = ImageLoaderHelper(path=path)

    model, device = ModelLoadHelper(weights='../yolov5/weights/yolov5s.pt', device='cpu')

    pbar = tqdm(loader)
    dt = [0.0, 0.0, 0.0]
    loss = []
    for batch_i, (img, targets, paths, shapes) in enumerate(pbar):
        # 图像张量预处理
        t1 = time_sync()
        img = img.to(device).float()
        targets = targets.to(device)
        img /= 255
        nb, _, height, width = img.shape
        t2 = time_sync()
        dt[0] += t2 - t1

        # 推理
        out, train_out = model(img)
        dt[1] += time_sync() - t2

        # NMS

        # loss
        # t = compute_loss([x for x in train_out], labels, model.model)
        # loss.append(t[1].numpy())

        # # Metrics
        # for si,pred in enumerate(out):
        #     lbs = targets

        # for i,det in enumerate(pred):
        #     p,img0,frame = path,img0


    with open('loss.txt', 'w') as f:
        for i in range(len(loss)):
            t = loss[i].tolist()

            for v in t:
                f.write(str(v))
                f.write('  ')

            f.write('\n')
